package com.xhs.hawaii.framework.thrift;

import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.imps.CuratorFrameworkState;
import org.apache.curator.framework.recipes.cache.ChildData;
import org.apache.curator.framework.recipes.cache.PathChildrenCache;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheEvent;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheListener;

import java.net.InetSocketAddress;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author tangchuang
 * @version 1.0
 * @since 2018/1/20
 */
@Slf4j
public class ThriftServerAddressContainer {

    private static final Integer DEFAULT_WEIGHT = 1;

    //服务地址集合
    private volatile List<InetSocketAddress> container;

    private CountDownLatch countDownLatch;

    private AtomicInteger index = new AtomicInteger(0);

    public ThriftServerAddressContainer(CuratorFramework zkClient, String path) {
        countDownLatch = new CountDownLatch(1);
        if (zkClient.getState() == CuratorFrameworkState.LATENT) {
            zkClient.start();
        }
        try {
            buildPathChildrenCache(zkClient, path, true);
        } catch (Exception e) {
            log.error("start zk server error", e);
            throw new RuntimeException("start zk server error");
        }
    }

    private void buildPathChildrenCache(final CuratorFramework client, String path, Boolean cacheData) throws Exception {
        PathChildrenCache cachedPath = new PathChildrenCache(client, path, cacheData);
        cachedPath.getListenable().addListener(new PathChildrenCacheListener() {
            @Override
            public void childEvent(CuratorFramework client, PathChildrenCacheEvent event) throws Exception {
                PathChildrenCacheEvent.Type eventType = event.getType();
                switch (eventType) {
                    case CONNECTION_RECONNECTED:
                        log.info("Connection is reconection.");
                        break;
                    case CONNECTION_SUSPENDED:
                        log.info("Connection is suspended.");
                        break;
                    case CONNECTION_LOST:
                        log.warn("Connection error,waiting...");
                        return;
                    case INITIALIZED:
                        log.warn("Connection init ...");
                    default:
                }
                // 任何节点的时机数据变动,都会rebuild,此处为一个"简单的"做法.
                cachedPath.rebuild();
                rebuild();
                countDownLatch.countDown();
            }

            protected void rebuild() throws Exception {
                List<ChildData> children = cachedPath.getCurrentData();
                if (children == null || children.isEmpty()) {
                    container = Collections.EMPTY_LIST;
                    log.error("thrift rpc-cluster error....");
                    return;
                }
                List<InetSocketAddress> current = new ArrayList<>();
                String pathP;
                for (ChildData data : children) {
                    pathP = data.getPath();
                    pathP = pathP.substring(path.length() + 1);
                    String address = new String(pathP.getBytes(), "utf-8");
                    current.addAll(transfer(address));
                }
                Collections.shuffle(current);
                container = current;
                index.set(0);
            }
        });
    }

    private List<InetSocketAddress> transfer(String address) {
        String[] hostname = address.split(":");
        Integer weight = DEFAULT_WEIGHT;
        if (hostname.length == 3) {
            weight = Integer.valueOf(hostname[2]);
        }
        String ip = hostname[0];
        Integer port = Integer.valueOf(hostname[1]);
        List<InetSocketAddress> result = new ArrayList<>();
        // 根据优先级，将ip：port添加多次到地址集中，然后随机取地址实现负载
        for (int i = 0; i < weight; i++) {
            result.add(new InetSocketAddress(ip, port));
        }
        return result;
    }

    //轮训获取
    public InetSocketAddress selector() {
        while (true) {
            if (container.isEmpty()) {
                return null;
            }
            int current = index.get();
            int next = (current + 1) % container.size();
            if (index.compareAndSet(current, next)) {
                return container.get(next);
            }
        }
    }

}
